import os
import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from pycocotools.coco import COCO
from segment_anything.utils.transforms import ResizeLongestSide

class AARDataset(Dataset):
    def __init__(self, cfg):
        super().__init__()
        self.root_dir = cfg['dataset']['root_dir']
        self.img_path = os.path.join(self.root_dir, 'images')
        self.gt_path = os.path.join(self.root_dir, 'ground_truth')
        # 文件列表
        self.img_file_list = sorted(os.listdir(self.img_path))
        self.gt_file_list = sorted(os.listdir(self.gt_path))

        # 类别
        self.category = cfg['dataset']['category']
        self.category_dict = {cat: i for i, cat in enumerate(self.category)}

        # 数据列表
        self.data_list = []
        for img_file, gt_file in zip(self.img_file_list, self.gt_file_list):
            self.data_list.append({
                'img': os.path.join(self.img_path, img_file),
                'gt': os.path.join(self.gt_path, gt_file)
            })
        print('Dataset size:', len(self.data_list))

        # 输入尺寸
        self.target_size = cfg['dataset']['target_size']
        # 正负样本点数目
        self.sample_point_num = cfg['dataset']['sample_num']
    
    def __len__(self):
        return len(self.data_list)
    
    def __getitem__(self, idx):
        data_id = self.data_list[idx]
        f_id = data_id['img'].split('/')[-1].split('.')[0]
        img = cv2.imread(data_id['img'])
        img = cv2.resize(img, self.target_size)
        img = img / 255.0
        img = img.transpose(2, 0, 1)
        img = img.astype('float32')
        gt = cv2.imread(data_id['gt'], cv2.IMREAD_GRAYSCALE)
        gt = cv2.resize(gt, self.target_size)
        gt = gt / 255.0
        gt = gt.astype('float32')

        # 生成box样本
        bbox_data = self.get_boxes_from_mask(gt, max_pixel=0)[0]
        # 生成point提示
        point_coords, point_labels = self.get_points_from_mask(gt, max_pixel=0)

        return {
            'img': img,
            'gt': gt,
            'bbox_data': bbox_data,
            'point_coords': point_coords,
            'point_labels': point_labels,
            'f_id': f_id
        }
    
    def get_boxes_from_mask(self, mask, max_pixel=0):
        """
        从mask中生成box样本
        """
        boxes = []
        contours, _ = cv2.findContours((mask * 255).astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if max_pixel > 0 and w * h > max_pixel:
                continue
            boxes.append([x, y, x + w, y + h])
        return boxes
    
    def get_points_from_mask(self, mask, max_pixel=0):
        """
        从mask中生成点样本
        """
        points = []
        labels = []
        contours, _ = cv2.findContours((mask * 255).astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if max_pixel > 0 and w * h > max_pixel:
                continue
            points.append([x + w // 2, y + h // 2])
            labels.append(1)
        return points, labels
    

class COCODataset(Dataset):

    def __init__(self, root_dir, annotation_file, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.coco = COCO(annotation_file)
        self.image_ids = list(self.coco.imgs.keys())

        # Filter out image_ids without any annotations
        self.image_ids = [image_id for image_id in self.image_ids if len(self.coco.getAnnIds(imgIds=image_id)) > 0]

    def __len__(self):
        return len(self.image_ids)

    def __getitem__(self, idx):
        image_id = self.image_ids[idx]
        image_info = self.coco.loadImgs(image_id)[0]
        image_path = os.path.join(self.root_dir, image_info['file_name'])
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        ann_ids = self.coco.getAnnIds(imgIds=image_id)
        anns = self.coco.loadAnns(ann_ids)
        bboxes = []
        masks = []

        for ann in anns:
            x, y, w, h = ann['bbox']
            bboxes.append([x, y, x + w, y + h])
            mask = self.coco.annToMask(ann)
            masks.append(mask)

        if self.transform:
            image, masks, bboxes = self.transform(image, masks, np.array(bboxes))

        bboxes = np.stack(bboxes, axis=0)
        masks = np.stack(masks, axis=0)
        return image, torch.tensor(bboxes), torch.tensor(masks).float()


def collate_fn(batch):
    images, bboxes, masks = zip(*batch)
    images = torch.stack(images)
    return images, bboxes, masks


class ResizeAndPad:

    def __init__(self, target_size):
        self.target_size = target_size
        self.transform = ResizeLongestSide(target_size)
        self.to_tensor = transforms.ToTensor()

    def __call__(self, image, masks, bboxes):
        # Resize image and masks
        og_h, og_w, _ = image.shape
        image = self.transform.apply_image(image)
        masks = [torch.tensor(self.transform.apply_image(mask)) for mask in masks]
        image = self.to_tensor(image)

        # Pad image and masks to form a square
        _, h, w = image.shape
        max_dim = max(w, h)
        pad_w = (max_dim - w) // 2
        pad_h = (max_dim - h) // 2

        padding = (pad_w, pad_h, max_dim - w - pad_w, max_dim - h - pad_h)
        image = transforms.Pad(padding)(image)
        masks = [transforms.Pad(padding)(mask) for mask in masks]

        # Adjust bounding boxes
        bboxes = self.transform.apply_boxes(bboxes, (og_h, og_w))
        bboxes = [[bbox[0] + pad_w, bbox[1] + pad_h, bbox[2] + pad_w, bbox[3] + pad_h] for bbox in bboxes]

        return image, masks, bboxes


def load_datasets(cfg, img_size):
    transform = ResizeAndPad(img_size)
    train = COCODataset(root_dir=cfg['dataset']['train']['root_dir'],
                        annotation_file=cfg['dataset']['train']['annotation_file'],
                        transform=transform)
    val = COCODataset(root_dir=cfg['dataset']['val']['root_dir'],
                      annotation_file=cfg['dataset']['val']['annotation_file'],
                      transform=transform)
    train_dataloader = DataLoader(train,
                                  batch_size=cfg['batch_size'],
                                  shuffle=True,
                                  num_workers=cfg['num_workers'],
                                  collate_fn=collate_fn)
    val_dataloader = DataLoader(val,
                                batch_size=cfg['batch_size'],
                                shuffle=True,
                                num_workers=cfg['num_workers'],
                                collate_fn=collate_fn)
    return train_dataloader, val_dataloader